package encoder

import (
	"fmt"
	"github.com/pion/rtp"
	"github.com/yapingcat/gomedia/go-mpeg2"
	"gogs.smarteye.com/servere/gobvrtp/track"
)

// PSEncoder 打包为ps包->打包为rtp包
type PSEncoder struct {
	RTPEncoder

	CodecIdList []track.CodecId

	muxer       *mpeg2.PSMuxer
	cidStreamId map[track.CodecId]uint8
}

// Init initializes the encoder.
func (e *PSEncoder) Init() error {

	var getCodecId = func(codec track.CodecId) mpeg2.PS_STREAM_TYPE {
		cid := mpeg2.PS_STREAM_H264
		switch codec {
		case track.CodecIdH264:
			cid = mpeg2.PS_STREAM_H264
		case track.CodecIdH265:
			cid = mpeg2.PS_STREAM_H265
		case track.CodecIdAAC:
			cid = mpeg2.PS_STREAM_AAC
		case track.CodecIdG711A:
			cid = mpeg2.PS_STREAM_G711A
		case track.CodecIdG711U:
			cid = mpeg2.PS_STREAM_G711U
		}

		return cid
	}

	e.muxer = mpeg2.NewPsMuxer()
	e.cidStreamId = map[track.CodecId]uint8{}

	for _, codec := range e.CodecIdList {
		streamId := e.muxer.AddStream(getCodecId(codec))
		e.cidStreamId[codec] = streamId
	}

	return e.RTPEncoder.Init()
}

func (e *PSEncoder) Encode(codec track.CodecId, nalus []byte) ([]*rtp.Packet, error) {

	var packets []*rtp.Packet
	e.muxer.OnPacket = func(pkg []byte) {
		if pkts, err := e.RTPEncoder.Encode(pkg, 0); err == nil {
			packets = append(packets, pkts...)
		}
	}

	streamId, ok := e.cidStreamId[codec]
	if !ok {
		return nil, fmt.Errorf("no stream")
	}

	err := e.muxer.Write(streamId, nalus, 0, 0)
	if err != nil {
		return nil, err
	}

	return packets, nil
}
